/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 * <p>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p>
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.sui.bigdata.flink.sql.launcher;

import com.sui.bigdata.flink.sql.core.util.PluginUtil;
import com.sui.bigdata.flink.sql.launcher.yarn.JobParameter;
import com.sui.bigdata.flink.sql.launcher.yarn.YarnClusterConfiguration;
import com.sui.bigdata.flink.sql.launcher.yarn.YarnClusterDescriptor;
import org.apache.commons.io.Charsets;
import org.apache.commons.lang.StringUtils;
import org.apache.flink.client.deployment.StandaloneClusterDescriptor;
import org.apache.flink.client.deployment.StandaloneClusterId;
import org.apache.flink.client.program.ClusterClient;
import org.apache.flink.client.program.rest.RestClusterClient;
import org.apache.flink.configuration.ConfigConstants;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.GlobalConfiguration;
import org.apache.flink.configuration.JobManagerOptions;
import org.apache.flink.core.fs.FileSystem;
import org.apache.flink.runtime.akka.AkkaUtils;
import org.apache.flink.runtime.util.LeaderConnectionInfo;
import org.apache.flink.yarn.AbstractYarnClusterDescriptor;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.yarn.api.records.ApplicationId;
import org.apache.hadoop.yarn.client.api.YarnClient;
import org.apache.hadoop.yarn.client.api.YarnClientApplication;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.apache.hadoop.yarn.exceptions.YarnException;
import org.apache.log4j.Logger;

import java.io.File;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.URLDecoder;
import java.util.*;
import java.util.stream.Collectors;
import java.util.stream.Stream;

import static java.util.Objects.requireNonNull;

/**
 * The Factory of ClusterClient
 *
 * @author
 */
public class ClusterClientFactory {
    private static Logger logger = Logger.getLogger(ClusterClientFactory.class);


    public static ClusterClient createStandaloneClient(LauncherOptions launcherOptions) throws Exception {
        String flinkConfDir = launcherOptions.getFlinkconf();
        Configuration config = GlobalConfiguration.loadConfiguration(flinkConfDir);

        StandaloneClusterDescriptor standaloneClusterDescriptor = new StandaloneClusterDescriptor(config);
        RestClusterClient clusterClient = standaloneClusterDescriptor.retrieve(StandaloneClusterId.getInstance());

        LeaderConnectionInfo connectionInfo = clusterClient.getClusterConnectionInfo();
        InetSocketAddress address = AkkaUtils.getInetSocketAddressFromAkkaURL(connectionInfo.getAddress());
        config.setString(JobManagerOptions.ADDRESS, address.getAddress().getHostName());
        config.setInteger(JobManagerOptions.PORT, address.getPort());
        clusterClient.setDetached(true);
        return clusterClient;
    }


    public static ClusterClient createYarnClient(String appName, String flinkPath, String yarnConfDir, String confProp, String udfJarPath) throws Exception {
        String flinkConfDir = flinkPath + "/conf";
        String flinkLib = flinkPath + "/lib";
        Configuration flinkConf = GlobalConfiguration.loadConfiguration(flinkConfDir);
        YarnConfiguration yarnConf;
        if (StringUtils.isNotBlank(yarnConfDir)) {
            try {
                flinkConf.setString(ConfigConstants.PATH_HADOOP_CONFIG, yarnConfDir);
                FileSystem.initialize(flinkConf);

                File dir = new File(yarnConfDir);
                ClusterClient clusterClient = null;

                if (dir.exists() && dir.isDirectory()) {
                    ApplicationId applicationId = null;

                    yarnConf = loadYarnConfiguration(yarnConfDir);
                    YarnClient yarnClient = YarnClient.createYarnClient();
                    haYarnConf(yarnConf);
                    yarnClient.init(yarnConf);
                    yarnClient.start();

                    confProp = URLDecoder.decode(confProp, Charsets.UTF_8.toString());

                    Properties confProperties = PluginUtil.jsonStrToObject(confProp, Properties.class);
                    System.out.println("confProp=" + confProperties);

                    applicationId = createApplication(yarnClient);

                    YarnClusterConfiguration clusterConf = getYarnClusterConfiguration(flinkConf, yarnConf, flinkConfDir);
                    //jobmanager+taskmanager param
                    JobParameter appConf = new JobParameter(confProperties);
                    System.out.println("wc :" + udfJarPath);
                    List<Path> userProvidedJars = new ArrayList();

                    File udfFiles = new File(udfJarPath);
                    for (File file : udfFiles.listFiles()) {
                        if (file.getAbsolutePath().contains(".jar") && !file.getAbsolutePath().endsWith("flink-stream-sql.jar")) {
                            userProvidedJars.add(new Path(file.toURI()));
                        }
                    }


                    YarnClusterDescriptor clusterDescriptor = new YarnClusterDescriptor(
                            clusterConf, yarnClient, appConf, applicationId, appName, userProvidedJars);
                    clusterClient = clusterDescriptor.deploy();

                    System.out.println("applicationId=" + applicationId.toString() + " has deploy!");
                    clusterClient.setDetached(true);
                    yarnClient.stop();
                }

                return clusterClient;
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }
        throw new UnsupportedOperationException("Haven't been developed yet!");
    }

    private static ApplicationId createApplication(YarnClient yarnClient) throws IOException, YarnException {
        YarnClientApplication app = yarnClient.createApplication();
        return app.getApplicationSubmissionContext().getApplicationId();
    }

    public static YarnClusterConfiguration getYarnClusterConfiguration(Configuration flinkConf, YarnConfiguration yarnConf, String flinkConfDir) {
        Path flinkJar = new Path(getFlinkJarFile(flinkConfDir).toURI());
        @SuppressWarnings("ConstantConditions") final Set<Path> resourcesToLocalize = Stream
                .of("logback.xml", "creep.xml")
                .map(x -> new Path(new File(flinkConfDir, x).toURI()))
                .collect(Collectors.toSet());
        String nameservices = yarnConf.get("dfs.nameservices", "10.201.7.127:8020");
        String appROOtDir = "hdfs://" + nameservices + "/user/data_team/flink/resource";
        System.out.println(" appROOtDir " + appROOtDir);
        return new YarnClusterConfiguration(flinkConf, yarnConf, appROOtDir, flinkJar, resourcesToLocalize);
    }

    public static final String FLINK_DIST = "flink-dist";

    private static File getFlinkJarFile(String flinkConfDir) {
        String errorMessage = "error not search " + FLINK_DIST + "*.jar";
        File[] files = requireNonNull(new File(flinkConfDir, "/../lib").listFiles(), errorMessage);
        Optional<File> file = Arrays.stream(files)
                .filter(f -> f.getName().startsWith(FLINK_DIST)).findFirst();
        return file.orElseThrow(() -> new IllegalArgumentException(errorMessage));
    }

    public static ClusterClient createYarnClient(String flinkConfDir, String yarnConfDir, String appId) {
        Long clusterTimestamp = Long.valueOf(appId.split("_")[1]);
        int id = Integer.valueOf(appId.split("_")[2]);
        Configuration config = GlobalConfiguration.loadConfiguration(flinkConfDir);
        if (StringUtils.isNotBlank(yarnConfDir)) {
            try {
                config.setString(ConfigConstants.PATH_HADOOP_CONFIG, yarnConfDir);
                FileSystem.initialize(config);

                YarnConfiguration yarnConf = YarnConfLoader.getYarnConf(yarnConfDir);
                YarnClient yarnClient = YarnClient.createYarnClient();
                yarnClient.init(yarnConf);
                yarnClient.start();
                ApplicationId applicationId = ApplicationId.newInstance(clusterTimestamp, id);

                AbstractYarnClusterDescriptor clusterDescriptor = new org.apache.flink.yarn.YarnClusterDescriptor(config, yarnConf, ".", yarnClient, false);
                ClusterClient clusterClient = clusterDescriptor.retrieve(applicationId);
                clusterClient.setDetached(true);
                yarnClient.stop();
                return clusterClient;
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }

        throw new UnsupportedOperationException("Haven't been developed yet!");
    }

    public static void killApp(String yarnConfDir, String appId) {
        Long clusterTimestamp = Long.valueOf(appId.split("_")[1]);
        int id = Integer.valueOf(appId.split("_")[2]);
        ApplicationId applicationId = ApplicationId.newInstance(clusterTimestamp, id);

        if (StringUtils.isNotBlank(yarnConfDir)) {
            try {
                YarnConfiguration yarnConf = YarnConfLoader.getYarnConf(yarnConfDir);
                YarnClient yarnClient = YarnClient.createYarnClient();
                yarnClient.init(yarnConf);
                yarnClient.start();
                yarnClient.killApplication(applicationId);
                logger.info(" kill " + applicationId.getId() + " sucessed");
                yarnClient.stop();
                cleanResource(yarnConf, appId);
            } catch (Exception e) {
                logger.info(" kill " + applicationId.getId() + " fail",e);
            }
        }
    }

    public static YarnConfiguration loadYarnConfiguration(String yarnConfDir) {
        org.apache.hadoop.conf.Configuration hadoopConf = new org.apache.hadoop.conf.Configuration();
        hadoopConf.set("fs.hdfs.impl", "org.apache.hadoop.hdfs.DistributedFileSystem");

        Stream.of("yarn-site.xml", "core-site.xml", "hdfs-site.xml").forEach(file -> {
            File site = new File(requireNonNull(yarnConfDir, "ENV HADOOP_CONF_DIR is not setting"), file);
            if (site.exists() && site.isFile()) {
                hadoopConf.addResource(new org.apache.hadoop.fs.Path(site.toURI()));
            } else {
                throw new RuntimeException(site + " not exists");
            }
        });

        YarnConfiguration yarnConf = new YarnConfiguration(hadoopConf);
        //        try (PrintWriter pw = new PrintWriter(new FileWriter(yarnSite))) { //write local file
        //            yarnConf.writeXml(pw);
        //        }
        return yarnConf;
    }

    /**
     * 处理yarn HA的配置项
     */
    private static org.apache.hadoop.conf.Configuration haYarnConf(org.apache.hadoop.conf.Configuration yarnConf) {
        Iterator<Map.Entry<String, String>> iterator = yarnConf.iterator();
        while (iterator.hasNext()) {
            Map.Entry<String, String> entry = iterator.next();
            String key = entry.getKey();
            String value = entry.getValue();
            if (key.startsWith("yarn.resourcemanager.hostname.")) {
                String rm = key.substring("yarn.resourcemanager.hostname.".length());
                String addressKey = "yarn.resourcemanager.address." + rm;
                if (yarnConf.get(addressKey) == null) {
                    yarnConf.set(addressKey, value + ":" + YarnConfiguration.DEFAULT_RM_PORT);
                }
            }
        }
        return yarnConf;
    }

    public static void cleanResource(String yarnConfDir) throws IOException {
        YarnConfiguration yarnConf = loadYarnConfiguration(yarnConfDir);
        String nameservices = yarnConf.get("dfs.nameservices", "10.201.7.127:8020");
        org.apache.hadoop.fs.FileSystem hdfs = org.apache.hadoop.fs.FileSystem.get(loadYarnConfiguration(yarnConfDir));
        Path appROOtDir = new Path("hdfs://" + nameservices + "/user/data_team/flink/resource");
        try {
            hdfs.delete(appROOtDir, true);
            hdfs.mkdirs(appROOtDir);
        } catch (IOException e) {
            logger.error(" Clean resource fail,Maybe next time. ");
        } finally {
            if (!hdfs.exists(appROOtDir)) {
                hdfs.mkdirs(appROOtDir);
            }
            hdfs.close();
        }

    }

    public static void cleanResource(YarnConfiguration yarnConf, String appId) {
        org.apache.hadoop.fs.FileSystem hdfs = null;
        try {
            String nameservices = yarnConf.get("dfs.nameservices", "10.201.7.127:8020");
            hdfs = org.apache.hadoop.fs.FileSystem.get(yarnConf);
            Path appROOtDir = new Path("hdfs://" + nameservices + "/user/data_team/flink/resource/" + appId);
            hdfs.delete(appROOtDir, true);
        } catch (Exception e) {
            logger.error("Clean resource " + appId + " delete file  fail.",e);
        } finally {
            try {
                hdfs.close();
            } catch (Exception e) {
                logger.error("Clean resource " + appId + " close hdfs fail.",e);

            }
        }
    }

}
